import time
import torch

# 定义一个函数来测量时间
def measure_time(device, X):
    X = X.to(device)
    time_start = time.time()
    Z = torch.mm(X, X)
    time_end = time.time()
    return round((time_end - time_start) * 1000, 2), device

# 检查CUDA是否可用
if torch.cuda.is_available():
    gpu_device = 'cuda:0'
else:
    print("CUDA is not available. Using CPU instead.")
    # gpu_device = 'cpu'

# 计算量较大的任务
X_large = torch.rand((10000, 10000))
cpu_time_large, cpu_device = measure_time('cpu', X_large)
gpu_time_large, gpu_device = measure_time(gpu_device, X_large)
print(f'Large task - CPU time cost: {cpu_time_large}ms on {cpu_device}')
print(f'Large task - GPU time cost: {gpu_time_large}ms on {gpu_device}')

# 计算量很小的任务
X_small = torch.rand((5000, 5000))
cpu_time_small, cpu_device = measure_time('cpu', X_small)
gpu_time_small, gpu_device = measure_time(gpu_device, X_small)
print(f'Small task - CPU time cost: {cpu_time_small}ms on {cpu_device}')
print(f'Small task - GPU time cost: {gpu_time_small}ms on {gpu_device}')
